/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2024-2024. All rights reserved.
 * Description:hccl comm 层RPC对外接口(NPU和CPU都包含)
 * Author: baochangjin
 * Create: 2024-02-25
 */
 
#ifndef HCCL_COMM_RPC_H
#define HCCL_COMM_RPC_H

#include "../common/hccl_comm_base.h"

namespace hccl {

class HcclCommRpc : public HcclCommBase {
public:

    explicit HcclCommRpc(u64 inCCLbufferSize = 0, u64 outCCLbufferSize = 0, std::string identifier = "");

    ~HcclCommRpc();

    HcclResult CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
        const std::vector<u32> &groupRanks, std::shared_ptr<HcclCommBase> &groupComm) override;

    HcclResult CreateGroup(const std::string &group, const u32 &groupRank, const u32 &userRank,
        const std::vector<u32> &groupRanks, std::shared_ptr<hcclComm> &groupComm) override;
    HcclResult InitImpl(DevType deviceType) override;

    HcclResult CollRecvUpdateRequest(void* keys, int32_t keyCount, HcclDataType keyType, void* values,
        int32_t valueCount, HcclDataType valueType, int32_t tag,  ServiceHandle* handle, UpdateReqStatus* status);
    HcclResult IsendUpdateResponse(ServiceHandle handle, HcclRequest* request);

    HcclResult CollRecvLookupRequest(void* keys, u32 keyCount, HcclDataType keyType,
        s32 tag, ServiceHandle* handle, LookupReqStatus* status, bool isWaitAny = false);

    HcclResult IsendLookupResponse(void* values, u32 count, HcclDataType type, ServiceHandle handle,
        HcclRequest* request);

    HcclResult GetLookupRequest(s32 tag, ServiceHandle* handle, u64* keys, u64 keyMaxNum, ReqStatus* request);

    HcclResult IsetLookupResponse(void* values, u32 count, HcclDataType type, ServiceHandle handle,
        HcclRequest *request);

    HcclResult WaitLookupData(s32 tag, void *value, void *stream, s32 &outZerocpyFlag, HcclRequest request);

    HcclResult RemoteLookUp(s32 tag, s32 insertFlag, void *stream, s32 &intZerocpyFlag,
        s32 &outZerocpyFlag, RemoteLookupPubParams &lookupPubInfo, HcclRequest* request, u64 keyNum = 0);

    HcclResult CollRemoteUpdate(s32 tag, void* keys, u32 keyMaxNum, void *values, void *tableId, s32 insertOption,
        s32 valueItemSize, rtStream_t stream, s32 maxEmbeddingDim, u64 keyNum = 0);

    HcclResult CollRemoteUpdatePaired(const HcomRemoteOperationParams &params);

    HcclResult CollRemoteLookup(s32 tag, void* keys, u32 keyMaxNum, void *values, void *tableId, s32 insertOption,
        s32 valueItemSize, rtStream_t stream, s32 flags, s32 maxEmbeddingDim, u64 keyNum = 0);
    HcclResult CollRemoteLookUp(const HcomRemoteOperationParams &params);
    HcclResult CollRemoteLookupPaired(const HcomRemoteOperationParams &params);
    HcclResult CollRemoteLookupUniquedPaired(const HcomRemoteOperationParams &params);

    void CollRemotePairedParaAssign(HcclOperator<EmbeddingServiceParam> &para,
        const HcomRemoteOperationParams &params);
    HcclResult AbortSelf(s32 tag);
    HcclResult RegTransportLinks(s32 linkNum, void *transportPara);
    HcclResult PsGetBoardVersion(bool &isBoardVersion);
    HcclResult WorkerGetBoardVersion(bool &isBoardVersion);
    HcclResult WaitSome(s32 count, HcclRequest requestArray[], s32* compCount, s32 compIndices[],
        HcclStatus compStatus[]);

    HcclResult Isend(void *buffer, u64 count, HcclDataType dataType, u32 peerRank, s32 tag, HcclRequest &request,
        HcclUserRequire &userRequire) const;
    HcclResult Improbe(u32 peerRank, s32 tag, s32 &flag, HcclMessage &msgHandle, HcclStatus &status) const;
    HcclResult Imrecv(void *buffer, u64 count, HcclDataType dataType, HcclMessage msg, HcclRequest &request);
    HcclResult SetKeyNumAndValCount(HcclOperator<EmbeddingServiceParam> &para, u64 keyNum, u64 keyMaxNum,
        u64 valueCount);
    };
}
#endif